from datetime import datetime
import pyqtgraph as pg
from PyQt5.QtCore import QDateTime, QTimer, QThread, pyqtSignal, Qt
from PyQt5.QtWidgets import QVBoxLayout, QHBoxLayout, QLabel, QPushButton,\
    QWidget, QCheckBox, QLineEdit, QTextEdit, QMessageBox, QSplitter
from pysnmp.entity.engine import SnmpEngine
from pysnmp.hlapi import getCmd, CommunityData, UdpTransportTarget, ContextData
from pysnmp.smi.rfc1902 import ObjectType, ObjectIdentity
from common_config_param import GlobalParam
from common_function import ip_format_check, text_input_check, welcome_init, clear_log, log_format_print

global global_running_flag


def snmp_get(ip, community, oid):
    try:
        iterator = getCmd(SnmpEngine(),
                          CommunityData(community),
                          UdpTransportTarget((ip, 161), timeout=1.0, retries=1),
                          ContextData(),
                          ObjectType(ObjectIdentity(oid)))
        errorIndication, errorStatus, errorIndex, varBinds = next(iterator)
        if errorIndication:
            return {"code": "errorIndication", "res": errorIndication}
        elif errorStatus:
            return {"code": "errorStatus", "res": errorStatus}
        else:
            for varBind in varBinds:
                return {"code": "200", "res": int(varBind.prettyPrint().split('=')[1].strip())}
    except Exception as e:
        return {"code": "Exception", "res": e}


class ThreadExtendsQThread(QThread):  # 定义线程类
    _signal = pyqtSignal(dict)  # 定义带参数一个信号

    def __init__(self, ip, community_name):
        super().__init__()
        self.ip = ip
        self.community_name = community_name
        self.snmp_oid_map = {"cpuCurrentUtilRate": "1.3.6.1.4.1.26067.1.9.2.1.0",
                             "cpuLongTimeUtilRate": "1.3.6.1.4.1.26067.1.9.2.2.0",
                             "devMemoryTotalNum": "1.3.6.1.4.1.26067.1.9.2.3.0",
                             "devMemoryFreeNum": "1.3.6.1.4.1.26067.1.9.2.4.0",
                             "devMemoryAllocNum": "1.3.6.1.4.1.26067.1.9.2.5.0"}

    def run(self):  # 线程的执行方法
        res = {"code": "200", "res": {"cpuCurrentUtilRate": 0, "memCurrentUtilRate": 0}, "srcData": []}
        res_list = []
        for key, value in self.snmp_oid_map.items():
            temp = snmp_get(self.ip, self.community_name, value)
            if not temp['code'] == "200":
                self._signal.emit({"code": "404", "res": temp["res"]})
                return
            res_list.append(temp["res"])
        res["res"]["cpuCurrentUtilRate"] = res_list[0]
        res["res"]["memCurrentUtilRate"] = int(100 - (res_list[3] / res_list[2]) * 100)
        res["srcData"] = res_list
        self._signal.emit(res)


class CpuMemMonitorLayout(QVBoxLayout):
    def __init__(self):
        super().__init__()

        self.running_flag = False

        self.ip = "127.0.0.1"
        self.snmp_community_name = "public"

        self.start_time = QDateTime.currentDateTime()
        self.timer = QTimer(self)

        self.timer.start(3000)

        self.thread_do = ThreadExtendsQThread(self.ip, self.snmp_community_name)
        self.thread_do._signal.connect(self.update_plot)

        self.data_time = []
        self.data_cpu = []
        self.data_mem = []

        top_layout = QHBoxLayout()

        ip_label = QLabel('DUT IP Address:')
        self.ip_input = QLineEdit()
        self.ip_input.setPlaceholderText("Enter DUT IP address")
        self.ip_input.setText("127.0.0.1")
        self.ip_input.editingFinished.connect(self.validate_ip)

        snmp_version_label = QLabel('SNMP Version:')

        self.snmp_version_input = QLineEdit()
        self.snmp_version_input.setFixedWidth(50)
        self.snmp_version_input.setText("2c")
        self.snmp_version_input.setReadOnly(True)

        snmp_community_label = QLabel('SNMP Community Name:')
        self.snmp_community_input = QLineEdit()
        self.snmp_community_input.setText("public")
        self.snmp_community_input.editingFinished.connect(self.validate_snmp_community_name)

        self.start_button = QPushButton('Start')
        self.start_button.clicked.connect(self.start_test)
        self.start_button.setEnabled(True)

        self.stop_button = QPushButton('Stop')
        self.stop_button.clicked.connect(self.stop_test)
        self.stop_button.setEnabled(False)

        clear_log_button = QPushButton("Clear Log")
        clear_log_button.clicked.connect(lambda x: clear_log(self.log_area))

        self.export_button = QPushButton('Export Image/Log')
        self.export_button.clicked.connect(self.export_data)

        top_layout.addWidget(ip_label)
        top_layout.addWidget(self.ip_input)
        top_layout.addWidget(snmp_version_label)
        top_layout.addWidget(self.snmp_version_input)
        top_layout.addWidget(snmp_community_label)
        top_layout.addWidget(self.snmp_community_input)
        top_layout.addWidget(self.start_button)
        top_layout.addWidget(self.stop_button)
        top_layout.addWidget(clear_log_button)
        top_layout.addWidget(self.export_button)

        self.addLayout(top_layout)

        # 中间布局
        middle_layout = QHBoxLayout()

        splitter = QSplitter(Qt.Horizontal)
        left_widget = QWidget()
        left_layout = QVBoxLayout(left_widget)

        plot_label = QLabel("Real-time Line Charts")

        self.plot_widget = pg.PlotWidget()
        self.plot_widget.setMenuEnabled(False)
        self.plot_widget.setLabel("left", "Rate(%)")
        self.plot_widget.setLabel("bottom", "Time(s)")

        self.plot_widget.addLegend()

        self.plot_cpu = self.plot_widget.plot(pen='g', name="CPU Utilization")
        self.plot_mem = self.plot_widget.plot(pen='y', name="Memory Utilization")

        left_layout.addWidget(plot_label)
        left_layout.addWidget(self.plot_widget)

        splitter.addWidget(left_widget)

        right_widget = QWidget()
        right_layout = QVBoxLayout(right_widget)

        self.cpu_checkbox = QCheckBox("CPU")
        self.cpu_checkbox.setChecked(True)
        self.cpu_checkbox.stateChanged.connect(self.update_plot_visibility)
        right_layout.addWidget(self.cpu_checkbox)

        self.mem_checkbox = QCheckBox("Memory")
        self.mem_checkbox.setChecked(True)
        self.mem_checkbox.stateChanged.connect(self.update_plot_visibility)

        right_layout.addWidget(self.mem_checkbox)

        right_layout.addStretch()

        splitter.addWidget(right_widget)

        splitter.setStretchFactor(0, 7)
        splitter.setStretchFactor(1, 3)

        mid_layout = QHBoxLayout()
        mid_layout.addWidget(splitter)

        middle_layout.addWidget(splitter)

        self.addLayout(middle_layout)

        log_label = QLabel("Real-time Logs")
        self.log_area = QTextEdit()
        self.log_area.setPlaceholderText("Log区域")
        self.log_area.setReadOnly(True)

        self.addWidget(log_label)
        self.addWidget(self.log_area)

        welcome_init(self.log_area)

    def validate_ip(self):
        if ip_format_check(self.ip_input.text()):
            QMessageBox.warning(self.parentWidget().parentWidget(),
                                "Invalid IP",
                                "The entered IP address is invalid.")

    def validate_snmp_community_name(self):
        if text_input_check(self.snmp_community_input.text()):
            QMessageBox.warning(self.parentWidget().parentWidget(),
                                "Invalid SNMP Community Name",
                                "The entered SNMP Community Name is invalid.")

    def start_test(self):
        if GlobalParam.global_running_flag:
            QMessageBox.warning(self.parentWidget().parentWidget(), "WARNING",
                                "Other programs are running, please stop first")
            return
        self.ip = self.ip_input.text()
        self.snmp_community_name = self.snmp_community_input.text()
        if ip_format_check(self.ip) or text_input_check(self.snmp_community_name):
            return
        GlobalParam.global_running_flag = True
        log_format_print(self.log_area,
                         "DUT Config:[IP:" + self.ip + "] - [SNMP Version:" +
                         self.snmp_version_input.text() + "] - [SNMP Community Name:" +
                         self.snmp_community_name + "]")
        QMessageBox.warning(self.parentWidget().parentWidget(),
                            "Attention",
                            "After starting a new monitoring, the old data will be erased")
        self.data_time = []
        self.data_cpu = []
        self.data_mem = []
        log_format_print(self.log_area, "CPU and Memory Utilization Monitor ---- Started")
        self.start_button.setEnabled(False)
        self.stop_button.setEnabled(True)
        self.running_flag = not self.running_flag
        self.start_time = QDateTime.currentDateTime()

        self.thread_do.ip = self.ip
        self.thread_do.community_name = self.snmp_community_name
        self.thread_do.start()
        self.timer.timeout.connect(self.timeout_handel)

    def timeout_handel(self):
        if not self.running_flag:
            return
        self.thread_do.ip = self.ip
        self.thread_do.community_name = self.snmp_community_name
        self.thread_do.start()

    def stop_test(self):
        self.running_flag = not self.running_flag
        self.start_button.setEnabled(True)
        self.stop_button.setEnabled(False)
        self.timer.timeout.disconnect()
        GlobalParam.global_running_flag = False
        log_format_print(self.log_area, "CPU and Memory Utilization Monitor ---- Stop")

    def update_plot(self, res):

        if not self.running_flag:
            return
        cpu_util = 0
        mem_util = 0
        if res["code"] == "404":
            log_format_print(self.log_area, res["res"])
        elif res["code"] == "200":
            cpu_util = res["res"]["cpuCurrentUtilRate"]
            mem_util = res["res"]["memCurrentUtilRate"]
            log_format_print(self.log_area, res["srcData"])

        if cpu_util >= 80:
            log_format_print(self.log_area, "CPU utilization is currently over 80% in use")
        if mem_util >= 80:
            log_format_print(self.log_area, "Memory utilization is currently over 80% in use")
        current_time = (QDateTime.currentDateTime().toMSecsSinceEpoch() - self.start_time.toMSecsSinceEpoch()) / 1000

        self.data_time.append(current_time)
        self.data_cpu.append(cpu_util)
        self.data_mem.append(mem_util)
        self.plot_cpu.setData(self.data_time, self.data_cpu)
        self.plot_mem.setData(self.data_time, self.data_mem)
        self.update_plot_visibility()

    def update_plot_visibility(self):
        if self.cpu_checkbox.isChecked():
            self.plot_cpu.setData(self.data_time, self.data_cpu)
        else:
            self.plot_cpu.clear()

        if self.mem_checkbox.isChecked():
            self.plot_mem.setData(self.data_time, self.data_mem)
        else:
            self.plot_mem.clear()

    def export_data(self):
        if self.running_flag:
            QMessageBox.warning(self.parentWidget().parentWidget(),
                                "Attention",
                                "Monitoring, stop monitoring and export!")
            return
        file_name_pre = str(datetime.now().strftime("%Y-%m-%d_%H%M%S_IP_")) + self.ip_input.text()

        exporter = pg.exporters.ImageExporter(self.plot_widget.plotItem)
        exporter.parameters()['width'] = 1000
        file_name = file_name_pre + "_cpu_mem_utilization_plot.png"
        exporter.export(file_name)

        log_format_print(self.log_area, "The image was successfully exported. File Name:[" + file_name + "]")

        file_name = file_name_pre + "_cpu_mem_utilization_monitor.log"

        with open(file_name, 'w') as f:
            f.write(self.log_area.toPlainText())

        log_format_print(self.log_area, "The logs were successfully exported. File Name:[" + file_name + "]")
